# -*- coding: UTF-8 -*-
# cython: language_level=3
import time
from concurrent import futures

from tqdm import tqdm

from base.base_parser import fenlei
from db.mongodb import MongoDB
from utils.get_domain_info import get_website_info
from utils.ip_area import find_ip
from utils.own_tools import get_current_date


def save_icp(oid, host, home_page='', title='', ip='', icp='', video_license='', visit_status=0, is_jump=False, *args,
             **kwargs):
    db = MongoDB()
    try:
        assert ip
        country, province_name, city_name, isp = find_ip(ip)
        if country and country != '中国':
            province_name = country + province_name
    except:
        province_name = ''
        city_name = ''

    if host.endswith('.gov.cn') or host.endswith('.edu.cn') or host.endswith('.gov') or host.endswith('.edu'):
        status = 8
    else:
        if visit_status == 1:
            status = 3
        else:
            status = fenlei(title, icp=icp, province_name=province_name)

    title = str(title or '').replace("'", '')[:100]

    data = {'home_page': home_page, 'icp': icp, 'PROVINCE_NAME': province_name, 'CITY_NAME': city_name, 'title': title,
            'license': video_license, 'IP': ip, 'status': status}
    if visit_status == 0:
        data['last_visit'] = time.strftime('%y%m%d')

    if is_jump:
        data.update({'host': host, 'record_time': get_current_date()})
        db.add('WEB_urls', data)
        db.update('WEB_urls', {'_id': oid}, {'status': 3, 'title': '302'})
    else:
        db.update('WEB_urls', {'_id': oid}, data)


def add_root_url(home_page, host, _id):
    dict_website_info = get_website_info(home_page, host=host)
    is_jump = (host != dict_website_info['host'])

    save_icp(_id, is_jump=is_jump, **dict_website_info)


def check(INPUT_URLS, n=None):
    len_urls = len(INPUT_URLS)
    # print('添加url数量： ', len_urls)
    p_num = min(len_urls, 40)
    if not p_num:
        return

    tasks, results = [], []
    with futures.ThreadPoolExecutor(max_workers=p_num) as executor:
        for i in tqdm(INPUT_URLS, total=len_urls, desc='add_task'):
            try:
                tasks.append(executor.submit(add_root_url, i['home_page'], i['host'], i['_id']))
            except:
                pass
        for task in tqdm(futures.as_completed(tasks), total=len_urls, desc=f'update_mongo-{n}'):
            # results.append(task.result())
            pass


def check_main(max_workers=10):
    db = MongoDB()
    # INPUT_URLS = db.find('WEB_urls', {"last_visit": {'$exists': True}}, sort=[('_id', -1)])
    INPUT_URLS = db.find('WEB_urls', {
        "status": {"$nin": [0, 1, 2, 7, 8, 9, 10]},
        "last_visit": {'$ne': time.strftime('%y%m%d')},
    }, limit=100000)
    # INPUT_URLS = db.find('WEB_urls', {})
    # INPUT_URLS = db.find('WEB_urls', {"status": 0})
    # INPUT_URLS = db.find('WEB_urls', {"status": 3})

    len_urls = len(INPUT_URLS)
    print('添加url数量： ', len_urls)

    if not len_urls:
        return True
    step = int(len_urls / max_workers) + 1
    INPUT_URLS = [INPUT_URLS[i:i + step] for i in range(0, len_urls, step)]

    tasks, results = [], []
    with futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
        for i, data in enumerate(INPUT_URLS):
            tasks.append(executor.submit(check, data, i))
        for task in tqdm(futures.as_completed(tasks), total=max_workers, desc='总进度'):
            # results.append(task.result())
            pass
    # print(results)


def test():
    db = MongoDB()
    INPUT_URLS = db.find('WEB_urls', {"status": 0}, sort=[('_id', -1)], limit=1000)
    for i in INPUT_URLS:
        # try:
        add_root_url(i['home_page'], i['host'], i['_id'])
    # except:
    #     pass


if __name__ == "__main__":
    check_main(max_workers=3)
    # test()
